import cv2
import numpy as np

import Utils

"""
A module for cell segmentation 
This document contains preprocess, threshold selecting, watershed for segmentation,
calculate and draw contour with different color.
"""


def countThreshold(img):
    flatArr = img.flatten()
    num = np.argmax(np.bincount(flatArr))
    return num


def process(img):
    gray = np.zeros_like(img)
    cv2.normalize(img, gray, 0, 255, cv2.NORM_MINMAX)
    ret, image = cv2.threshold(gray, countThreshold(gray), 255, cv2.THRESH_BINARY)

    return ret, image


def threshold(img):
    # 1. simplest threshold
    hist, bins = np.histogram(img.flatten(), 256, [0, 256])
    th = np.where(hist == np.max(hist))[0]
    # print(th[0])
    ret1, th1 = cv2.threshold(img, th[0] + 1, 255, cv2.THRESH_BINARY)
    # ret1, th1 = cv2.threshold(image, 150 , 255, cv2.THRESH_BINARY)
    # 2. adaptive threshold MEAN
    th2 = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 3, 2)
    # 3. GAUSSIAN threshold
    th3 = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 3, 2)

    # 4. using OTUS
    # Otsu's thresholding
    ret4, th4 = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # plot all the images and their histograms
    images = [img, 0, th1,
              img, 0, th2,
              img, 0, th3,
              img, 0, th4]

    # just for test
    # titles = ['Original Image', 'Histogram', 'Global Thresholding (v=th[0]+1)',
    #           'Original Image', 'Histogram', "Adaptive Mean Thresholding",
    #           'Original Image', 'Histogram', "Adaptive Gaussian Thresholding",
    #           'Original Image', 'Histogram', "Otsu's Thresholding"]
    #
    # for i in range(4):
    #     plt.subplot(4, 3, i * 3 + 1), plt.imshow(images[i * 3], 'gray')
    #     plt.title(titles[i * 3]), plt.xticks([]), plt.yticks([])
    #     plt.subplot(4, 3, i * 3 + 2), plt.hist(images[i * 3].ravel(), 256)
    #     plt.title(titles[i * 3 + 1]), plt.xticks([]), plt.yticks([])
    #     plt.subplot(4, 3, i * 3 + 3), plt.imshow(images[i * 3 + 2], 'gray')
    #     plt.title(titles[i * 3 + 2]), plt.xticks([]), plt.yticks([])

    # plt.show()

    # img = th4
    return th4


def contrast_stretching(img):
    I = img
    a = 0
    b = 255
    c = np.min(I)
    d = np.max(I)
    # contrast stretched
    O = ((I - c) * ((b - a) / (d - c))) + a
    return O.round().astype("uint8")


def pre_process_th(img):
    # img = cv2.equalizeHist(img)
    img = contrast_stretching(img)
    # img = cv2.medianBlur(img, 5)
    # thresholding(img)
    blur = cv2.GaussianBlur(img, (15, 15), 0)
    _, img = process(blur)
    return img


# watershed for segmentation
def watershed_process(img_raw, img_threshold):
    iterations_a, iterations_b = 2, 10
    dist_value = 5
    min_value, max_value = 0, 255
    factor = 0.25
    kernel = np.ones((3, 3), np.uint8)
    img_mor = cv2.morphologyEx(img_threshold, cv2.MORPH_OPEN, kernel, iterations=iterations_a)

    bg = cv2.dilate(img_mor, kernel, iterations=iterations_b)
    dist = cv2.distanceTransform(img_mor, cv2.DIST_L2, dist_value)
    _, sure_fg = cv2.threshold(dist, dist.max() * factor, max_value, cv2.THRESH_BINARY)

    flag = 1
    surface_fg = np.uint8(sure_fg)
    unknown = bg - surface_fg
    _, markers = cv2.connectedComponents(surface_fg)

    markers = markers + flag
    markers[unknown == 255] = min_value
    markers = cv2.watershed(img_raw, markers)

    return markers


def pre_process(img_color):
    img_new = img_color.copy()

    # contrast stretch
    img_gray = cv2.cvtColor(img_color, cv2.COLOR_BGR2GRAY)
    img_gray = contrast_stretching(img_gray)

    # _, img_threshold = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    _, img_threshold = process(img_gray)

    # find labels of watershed
    markers = watershed_process(img_color, img_threshold)

    # binarization
    img_new[markers == -1] = [255, 255, 255]
    img_ret = cv2.inRange(img_new, (255, 255, 255), (255, 255, 255)).astype(np.uint8)

    return img_ret, markers


def calc_contour(img, markers):
    contours_list = []
    for marker in np.unique(markers):
        if marker <= 1:
            continue
        mask = np.zeros(img.shape, dtype="uint8")
        mask[markers == marker] = 255
        contours = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contour = contours[1][0]
        area = cv2.contourArea(contour)
        if contour.shape[0] > 4 and area > 50:
            contours_list.append(contour)
    print(f'The current image contours: {len(contours_list)}')
    return contours_list


# abort
def calc_contours_(img):
    contours = []
    img_contour, contour, hierarchy = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    temp = contour[0]
    for c in contour:

        # By setting the minimum cell area, noises won't be detected as a cell
        if cv2.contourArea(c) > 20:
            contours.append(c)
    print(f'The current image contours: {len(contours)}')
    return img_contour, contours, hierarchy


# abort
def draw_contour_with_colors(img, contours, color_list, hierarchy):
    img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # Utils.show_img(img_rgb, "img_rgb of img_test")
    for i in range(0, len(contours)):
        image_draw = cv2.drawContours(img_rgb, contours, i, color_list[i], 2)

    # Utils.show_img(image_draw, "img_rgb after draw")
    return image_draw


def draw_cell_contour_and_ID(img, cell_list):
    image_draw = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    # Utils.show_img(img_rgb, "img_rgb of img_test")
    for cell in cell_list:
        image_draw = cv2.drawContours(image_draw, [cell.contour], -1, cell.segmentation_color, 2)
        text = str(cell.get_id())
        cv2.putText(image_draw, text, cell.get_centroid(), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1, 8, 0)
        if cell.is_splitting >=1:
            x, y, w, h = cell.get_bounding_rect()
            image_draw = cv2.rectangle(image_draw, (x, y), (x + w, y + h), (0, 0, 255), 2)

    # Utils.show_img(image_draw, "img_rgb after draw")
    return image_draw



